
{-

Note, this module also includes the compiler directive logic and
builtin processing.  

-}

module Compiler where

import Prelude hiding (lookup)
import EngineTypes
import Engine
import RPTypes
import SParser
import TableAccounting
import Query


import AtomMap hiding (lookupAtom)

import Data.Map

import Control.Monad
import Data.Array.IO
import Data.IORef
import Maybe
import Data.List (groupBy,sort,nub,maximumBy,find,map)
import System.Random
import Data.Int

import qualified Data.Set as Set

import Numeric


-- import Debug.Trace.Location
-- import Debug.Trace

type AMap = Map Atom AtomRef

--
-- QUERY/PROGRAM COMPILER
--

-- the WAMRef here is to permit the forward pointer to have its value 
-- changed to the compiled term
-- XXX think about this pointer - maybe there's a neater way

-- Compile a term directly onto the heap

{-
compile_heap_term :: (WAMData a) -> Term -> WAMRef -> ST a ()
compile_heap_term wd (V var) argref = 
                    do  ref <- getVMap wd var
                        if isJust $ ref
                           then do set_value' wd argref ( fromJust $ ref )
                                   return ()
                           else do set_variable' wd argref
                                   insertVMap wd var argref
                                   return ()

compile_heap_term wd (T at ar args) ref = do (put_structure' wd at ar ref)
                                             start <- getHT wd
                                             incHT wd ar
                                             compile_heap_term_args wd args start
                                             return ()

-- it should be possible to use a foldl-like mechanism here ...
compile_heap_term_args _ [] _ = return ()
compile_heap_term_args wd (a:as) st = do compile_heap_term wd a st
                                         compile_heap_term_args wd as (st + 1)

-}


-- Compile a clause into program unification instructions
compileSetup :: IO (WAMData a)
compileSetup = emptyWAM stacksize stacksize trailsize

compileQuery :: (WAMData a) -> CBody -> (IO (WAMRef,VarMap))
compileQuery wd cd = 
    do resetRegister wd
       emptyVMap wd
       start <- pushCode wd (Allocate (fromIntegral $ length (permRegs vmap)))
       compileQueryV wd cd vmap 
       -- pushCode wd Deallocate
       pushCode wd HaltOP
       return (start,vmap)
    where vmap = varMap (CLA (T "a" 0 [] nullpos) cd) -- a dummy head

compileQueryV :: (WAMData a) -> CBody -> VarMap -> (IO WAMRef)
compileQueryV wd c v =
    do start <- getCT wd
       amap <- compileAtoms wd [(CLA (T "a" 0 [] nullpos) c)]
       compile_body wd c v amap -- make the head a null
       return start

--compile :: [Clause] -> IO (WAMData a)
--compile c =  
--          do wd <- emptyWAM stacksize stacksize stacksize
--             compile_program wd c
--             return wd

-- predicate defined as a collection of clauses, with some hints added on
--  how to compile them

clauseID :: Predicate -> ClauseID
clauseID p = let (CLA h b) = head $ clauselist p
             in (CID (atom h) (arity h) 0)

compileProgram :: (WAMData a) -> [Clause] -> IO ()
compileProgram wd c =
    let directives = [x | x <- c, isDir x ]
        cls = [x | x <- c, not $ isDir x ]
        rawchunks = (groupBy (\x y -> (compare x y) == EQ) (sort cls))
        chunks = [ applyDirectives x directives | x <- rawchunks]
        randpreds = [predid x | x <- chunks, randomized x] 
    in do atommap <- compileAtoms wd c
          compileClausesets wd chunks atommap
          setRandoms wd randpreds
          applyGlobalDirs wd directives
          return ()


atomClause :: (WAMData a) -> Clause -> IO ()
atomClause wd (CLA h b) =
    do atomTerm wd h 
       atomCBody wd b
atomClause wd (DIR t) = atomTerm wd t

atomTerm wd (T at _ more _) =
    do Engine.newAtom wd at
       mapM (atomTerm wd) more
       return ()
atomTerm wd _ = return ()

atomCBody wd (CNIL) = return ()
atomCBody wd (CCONJ a b _) = do atomCBody wd a ; atomCBody wd b
atomCBody wd (CDISJ a b _) = do atomCBody wd a ; atomCBody wd b
atomCBody wd (CTERM t _) = atomTerm wd t

compileAtoms :: (WAMData a) -> [Clause] -> IO AMap
compileAtoms wd cs =
    do mapM (atomClause wd) cs
       am <- readIORef $ atomIndex wd
       return (revmap am)

applyGlobalDirs :: (WAMData a) -> [Clause] -> IO ()
applyGlobalDirs wd ds = 
	do mapM (\d -> mapM (\f -> f wd d) globaldirs) (Data.List.map getdir ds)
	   return ()

type GlobalDirective a = (WAMData a) -> Term -> IO ()

globaldirs :: [GlobalDirective a]
globaldirs = [ applySingleArgIntDirective "histDepth" setDepth,
	       applySingleArgFloatDirective "lambda" setLambda,
	       applySingleArgFloatDirective "scale" setScale,
	       applySingleArgIntDirective "seed" (\wd val -> do {writeIORef (sgen wd) (mkStdGen val);writeIORef (seed wd) val;} ),
               applySingleArgIntDirective "indTrialDepth" (\wd val -> do { writeIORef (shortThreshold wd) (fromIntegral $val);}),
	       applySingleArgIntDirective "shortThreshold" (\wd val -> do { writeIORef (shortThreshold wd) (fromIntegral $val); writeIORef (restartThreshold wd) (fromIntegral $val);} ),
	       setGlobalQuery
             ]

{-
restartThreshold = 5000000000 :: Word64 
--restartThreshold = 30000 :: Word64 
minimumRestartThreshold = RPTypes.restartThreshold `div` 2
restartThresholdScale = 0.9
weightScalePower = 5 :: Double
weightAdd = 0 :: Double
lambda = 0.60::Double 
-}


setGlobalQuery :: GlobalDirective a
setGlobalQuery wd t = do (_,d,i) <- normalizeQuery t
                         setCompiledQuery wd (d,i)

setCompiledQuery wd (d,i) = do writeIORef (dedQuery wd) d
                               writeIORef (indQuery wd) i

getCompiledQuery :: (WAMData a) -> IO (CBody,CBody)
getCompiledQuery wd = do d <- readIORef (dedQuery wd)
                         i <- readIORef (indQuery wd)
                         return (d,i)

applySingleArgIntDirective :: String -> ((WAMData a) -> Int -> IO ()) -> GlobalDirective a
applySingleArgIntDirective str func wd term@(T _ 1 [(T arg 0 [] _)] _) =
    if (atom term) == str && (length res) == 1 && as == ""
       then do func wd ai 
               return ()
       else return ()
    where res = (readDec arg)
          [(ai,as)] = res
applySingleArgIntDirective _ _ _ _ = return ()

applySingleArgFloatDirective :: String -> ((WAMData a) -> Double -> IO ()) -> GlobalDirective a
applySingleArgFloatDirective str func wd term@(T _ 1 [(T arg 0 [] _)] _) =
    if (atom term) == str && (length res) == 1 && as == ""
       then do func wd ai 
               return ()
       else return ()
    where res = (readFloat arg)
          [(ai,as)] = res
applySingleArgFloatDirective _ _ _ _ = return ()

-- parseTermA _ = (T "a" 0 [])

applyDirectives :: [Clause] -> [Clause] -> Predicate
applyDirectives cl directives 
     | regularatom pn && -- to deal with atoms like '=' which the parser cannot deal with
       elem (stringToTerm $ "randomize(" ++ pn ++ ")" ) [ getdir x | x <- directives] 
                   = PRED {predid = pid cl,
                           clauselist = cl,
                           randomized = True,
                           uwrandom = False}
     | regularatom pn && -- to deal with atoms like '=' which the parser cannot deal with
       elem (stringToTerm $ "uw_randomize(" ++ pn ++ ")" ) [ getdir x | x <- directives] 
                   = PRED {predid = pid cl,
                           clauselist = cl,
                           randomized = False,
                           uwrandom = True}
     | otherwise = PRED {predid = pid cl,
                         clauselist = cl,
                         randomized = False,
                         uwrandom = False}
   where pn = predicateName cl
         regularatom pn = isNothing (find (\x -> not (x `elem` namechars)) pn)
         predicateName ((CLA (T na _ _ _) _):_) = na
         pid ((CLA h _):_) = predID h

getdir (DIR x) = x

isDir (DIR _) = True
isDir _ = False

compileClausesets :: (WAMData a) -> [Predicate] -> AMap -> IO ()
compileClausesets _ [] _ = return ()
compileClausesets wd (c:cs) amap = 
    do if ((not $ randomized c) && (not $ uwrandom c))
         then compileClauseset wd (clauselist c) (clauseID c) amap
         else compileRandomClauseset wd (clauselist c) (clauseID c) (randomized c) amap
       compileClausesets wd cs amap

compileRandomClauseset :: (WAMData a) -> [Clause] -> ClauseID -> Bool -> AMap -> IO ()
compileRandomClauseset wd [c] _ _ amap = 
    do ddout "WARNING: Random predicate with 1 clause!\n"
       compileClauseset wd [c] NullCID amap
compileRandomClauseset wd c cid weighted amap =
    do (addrs,_,_) <- ccs wd c amap
       bot <- if weighted
                then pushCode wd (Randomize_try (fromIntegral (length addrs)))
                else pushCode wd (Randomize_unweighted_try (fromIntegral $ length addrs))
       ir1 wd addrs size cid 1
       insertCallIndex wd (findWithDefault 0 (atom h) amap) (arity h) bot
    where
       (x:_) = c
       (CLA h _) = x
       size = sizeHead (head c)

-- A bit hacky the next bit ...
ir1 wd (a:as) size cid n = 
   do ret <- pushCode wd (Try_there (a) size 0 (setCnum cid n))
      setCode wd ret (Try_there (a) size (ret) (setCnum cid n))
      ir2 wd (as) cid (n+1)
      
ir2 wd (a:as) cid n =
   if (length (a:as)) > 1
   then do ret <- pushCode wd (Retry_there (a) 0 (setCnum cid n))
           setCode wd ret (Retry_there (a) (ret) (setCnum cid n))
           ir2 wd (as) cid (n+1)
   else pushCode wd (Trust_there (a) (setCnum cid n))

-- the offset business below puts the index one back for complex clauses
-- to make the index match up with the try instruction

-- singleton clauses need no choicepoint

compileClauseset :: (WAMData a) -> [Clause] -> ClauseID -> AMap -> IO ()
compileClauseset wd [c] cid amap =  
    do (op,addr) <- compileClause wd c amap
       op addr
       return ()
compileClauseset wd c cid amap = 
    do (addrs,op,addr) <- ccs wd c amap
       ia1 wd addrs size cid 1
       op addr
    where size = sizeHead (head c)
	
sizeHead (CLA (T _ ar _ _) _) = ar

ccs wd (c:cs) amap = 
    do r <- pushCode wd NullOP; 
       (op,addr) <- compileClause wd c amap 
       rs <- ccs1 wd cs amap
       return ((r:rs),op,r)  	
	   
ccs1 _ [] amap = return []
ccs1 wd (c:cs) amap 
    = do r <- pushCode wd NullOP; 
         compileClause wd c amap
         rs <- ccs1 wd cs amap 
         return (r:rs)

setCnum (CID at ar _) n = (CID at ar n)

-- Try ; Retry; ... Retry; Trust
ia1 wd (a:(ab:as)) size cid n = 
   do setCode wd a (Try_here size (ab -1) (setCnum cid n))
      ia2 wd (ab:as) cid (n+1)
      
ia2 wd (a:(ab:as)) cid n =
   do setCode wd a (Retry_here (ab - 1) (setCnum cid n))
      ia2 wd (ab:as) cid (n+1)
ia2 wd [a] cid n = setCode wd a (Trust_here (setCnum cid n))

-- returned is code to generate an index for this clause
--  fixed disjunction in queries and clause bodies anyway ...
compileClause :: (WAMData a) -> Clause -> AMap -> IO (WAMRef -> IO (),WAMRef)
compileClause wd (CLA h CNIL) amap = 
    do emptyVMap wd
       resetRegister wd
       startaddr <- getCT wd
       dd ("COMPILING CLAUSE: " ++ (show (CLA h CNIL)) ++ "\n")
       dd ("vmap is " ++ (foldl (++) "" (Prelude.map show (assocs vmap))) ++ "\n")
       compileHead wd h vmap amap
       pushCode wd (Proceed)
       argsize <- getTopArgRegister wd  -- BROKEN IN HERE
       return (insertCallIndex wd (findWithDefault 0 (atom h) amap) (arity h), startaddr)
    where vmap = varMap (CLA h CNIL)

compileClause wd (CLA h t) amap = 
    do emptyVMap wd
       resetRegister wd
       dd ("COMPILING CLAUSE: " ++ (show (CLA h t)) ++ "\n")
       -- dd ("VMAP IS: " ++ (showTree vmap) ++ "\n")
       setTopTempRegister wd (fromIntegral (length $ tempRegs vmap))
       allocaddr <- pushCode wd (Allocate (fromIntegral $ length (permRegs vmap)))
       if (containsCut t) && not (isCutTerm $ leftmost t)
	  then pushCode wd (Get_level (findWithDefault NullRegister cutVariable vmap))
          else return 0
       compileHead wd h vmap amap
       if isCutTerm $ leftmost t
	  then pushCode wd Neck_cut
          else return 0
       compile_body wd (deleteLeftmostCut t) vmap amap
       -- regf <- getTopTempRegister wd
       -- setCode wd allocaddr (Allocate (length (permRegs vmap)))
       pushCode wd (Deallocate)
       argsize <- getTopArgRegister wd
       return (insertCallIndex wd (findWithDefault 0 (atom h) amap) (arity h), allocaddr)
    where vmap = varMap (CLA h t)

-- big awful ugly hack until compile_query is demonadified
compileHead wd h vmap amap = 
    do mapM (pushCode wd) ops
       mapM (\x -> newTempRegister vmap wd) (tempRegsC vm)
       mapM (\x -> newArgRegister wd) (argRegs vm)
       mapM (\(x,y) -> insertVMap wd x y) [(x,y) | (x,(y,n)) <- assocs vm, tr x, n>0 ]
    where (ops,vm) = (compileTermHack h vmap amap)
          tr x = (take 4 x) /= "ANON" -- no-one better call a variable "ANON"

tempRegs :: VarMap -> [Register]
tempRegs = elems . (Data.Map.filter fil)
           where fil (TempRegister x) = True
                 fil _ = False
 
permRegs :: VarMap -> [Register]
permRegs = elems . (Data.Map.filter fil)
           where fil (PermRegister x) = True
                 fil _ = False

varMap :: Clause -> VarMap
varMap (CLA h CNIL) = 
              fromList [ (vn x,TempRegister y) | (x,y) <- zip vars [0..(fromIntegral $length vars)] ]
              where vars = uniqueVars h
                    vn (V x _) = x 
varMap (CLA h b) =
        fromList ((regset TempRegister temps) ++ (regset PermRegister (perms `union` cutregs)))
        where body = mergeHead h b
              (temps,perms) = partition (\x -> x <= 1) counts
              counts = varCounts body
              regset typ m = [ (vn x,typ y) | (x,y) <- zip (keys m) [0::UnsignedNum .. ((fromIntegral $ length $ keys m)-1)] ]
              vn (V x _) = x
              cutregs = 
                        if (containsCut b) && not (isCutTerm $ leftmost b)
                           then  (fromList [(V cutVariable nullpos,2)]) :: (Map Term Int)
                           else  fromList []


varCounts (CTERM t _) = fromList [ (x,1) | x <- uniqueVars t ]
varCounts (CCONJ b1 b2 _) = unionWith (+) (varCounts b1) (varCounts b2)
varCounts (CDISJ b1 b2 _) = unionWith (+) (varCounts b1) (varCounts b2)
varCounts CNIL = Data.Map.empty

mergeHead :: Term -> CBody -> CBody
mergeHead (T _ _ args1 _) (CTERM (T n l args2 p) v) = (CTERM (T n l (args1 ++ args2) p) v)
mergeHead t (CCONJ l r v) = (CCONJ (mergeHead t l) r v)
mergeHead t (CDISJ l r v) = (CDISJ (mergeHead t l) r v)
mergeHead t CNIL = (CTERM t Notset)
mergeHead t (CTERM (CutTerm _) _) = (CTERM t Notset)

uniqueVars :: Term -> [Term]
uniqueVars = nub . var1

var1 :: Term -> [Term]
var1 (V x p) = [(V x p)]
var1 (T _ _ args _) = foldl (++) [] [var1 x | x <- args ]
var1 _ = []

bodlist (CCONJ x y _) = (x:bodlist y)
bodlist CNIL = []
bodlist x = [x]
 
-- compile_head wd (T _ _ args) vmap = compile_terms wd [(Nothing,x) | x <- args] vmap

compile_body :: (WAMData a) -> CBody -> VarMap -> AMap -> IO ()
compile_body wd (CTERM x _) vmap amap = 
    do compile_goal wd x vmap amap
       return ()
compile_body wd (CCONJ x y _) vmap amap = 
    do compile_body wd x vmap amap
       compile_body wd y vmap amap
compile_body wd CNIL _ _ = return ()    
compile_body wd (CDISJ x y _) vmap amap = 
    let disjchain = wrapupdisj (CDISJ x y Notset)
    in compiledisj wd disjchain vmap amap
    
compiledisj wd disjchain vmap amap = 
    do (tries,branches) <- cdj wd disjchain vmap amap
       inserttries wd tries
       finaladdr <- getCT wd
       insertbranches wd (Branch finaladdr) branches
       return ()

-- this is similar to the ccs code above - it would be nice to merge them ...

cdj _ [] _ _ = return ([],[])
cdj wd (h:hs) vmap amap = 
    do try <- pushCode wd NullOP
       compile_body wd h vmap amap
       branch <- pushCode wd (Branch 0)
       (triesrest,branchesrest) <- cdj wd hs vmap amap
       return (try:triesrest,branch:branchesrest)
       
inserttries wd (h:(h0:hs)) =
    do setCode wd h (Try_here 0 (h0 - 1) NullCID)
       it1 wd (h0:hs)
       
it1 wd [x] = 
    do setCode wd x (Trust_here NullCID)
       return ()
       
it1 wd (h:(h0:hs)) = 
    do setCode wd h (Retry_here (h0 - 1) NullCID)
       it1 wd (h0:hs)

-- last branch is redundant in compiled code and should be removed

insertbranches wd inst branches = mapM (\x -> setCode wd x inst) branches 
    
wrapupdisj (CDISJ x y _) = (x:(wrapupdisj y))
wrapupdisj CNIL = []
wrapupdisj x = [x]    

{-
This again will be try; retry; ... ; trust
-}

-- fails for neck cuts ...
compile_goal wd (CutTerm _) vmap _ = 
    do pushCode wd (Cut (findWithDefault NullRegister cutVariable vmap))
compile_goal wd (T at ar args _) vmap amap = 
    do resetArgRegister wd
       compile_query_list wd (zip [(True)|x <- args] args) vmap amap
       pushCode wd (Call (findWithDefault 0 at amap) ar)

-- Compile program tail terms/queries - returns the register of the term
--  the second argument is the indicator of whether we're at the top level

compile_query :: (WAMData a) -> Bool -> Term -> VarMap -> AMap -> IO Register
compile_query wd True (V var _) vmap amap =
    do varref <- getVMap wd var
       if (isNothing varref)
          then do arg_reg <- newArgRegister wd
                  -- norm_reg <- newTempRegister wd
                  pushCode wd (Put_variable norm_reg arg_reg)
                  insertVMap wd var norm_reg
                  return norm_reg
          else do arg_reg <- newArgRegister wd
                  pushCode wd (Put_value norm_reg arg_reg)
                  return (fromJust varref)
    where norm_reg = (findWithDefault NullRegister var vmap)
compile_query wd False (V var _) vmap amap =
    do varref <- getVMap wd var
       if (isNothing varref)
          then do -- reg <- newTempRegister wd
                  pushCode wd (Set_variable norm_reg)
                  insertVMap wd var norm_reg
                  return norm_reg
          else do pushCode wd (Set_value norm_reg)
                  return norm_reg
    where norm_reg = (findWithDefault NullRegister var vmap)
compile_query wd toplev (T at ar args _) vmap amap =
    do heapregs <- compile_query_list wd (zip [(False) | x <- args] args) vmap amap
       reg <- nReg wd
       spot <- pushCode wd (Put_structure (findWithDefault 0 at amap) ar reg)
       compile_query_args wd heapregs args vmap amap
       return reg
    where nReg = if toplev
                    then newArgRegister
                    else newTempRegister vmap


-- this bit first pushes the query bits on - variables are ignored
compile_query_list :: (WAMData a) -> [(Bool,Term)] -> VarMap -> AMap -> IO [Register]
compile_query_list _ [] _ _ = return []
compile_query_list wd ((False,V _ _):lts) vm am = compile_query_list wd lts vm am
compile_query_list wd ((True,V v p):lts) vm am =
    do compile_query wd True (V v p) vm am
       compile_query_list wd lts vm am
compile_query_list wd ((level,term):lts) vm am=
    do loc <- compile_query wd level term vm am
       locs <- compile_query_list wd lts vm am
       return (loc:locs)
       
compile_query_args :: (WAMData a) -> [Register] -> [Term] -> VarMap -> AMap -> IO ()
compile_query_args _ [] [] _ _ = return () 
compile_query_args wd rs ((V v p):ts) vm am =
    do compile_query wd False (V v p) vm am
       compile_query_args wd rs ts vm am
compile_query_args wd (r:rs) (_:ts) vm am =
    do pushCode wd (Set_value r)
       compile_query_args wd rs ts vm am

 
--- Compile program head terms
-- Bool == toplevel 

varMapToVarMapCount :: VarMap -> VarMapCount
varMapToVarMapCount x = fromList [ (v,(r,0)) | (v,r) <- assocs x ]

seenVariable :: Variable -> VarMapCount -> Bool
seenVariable v vmc =
	let lu = Data.Map.lookup v vmc
	    (r,c) = fromJust lu
	in (isJust lu) && (c > 0)

addSeen :: Variable -> VarMapCount -> VarMapCount
addSeen v vmc = adjust (\(x,y) -> (x,y+1)) v vmc

assignArgRegister :: Maybe Variable -> VarMapCount -> (Register,VarMapCount)
assignArgRegister vname vmc = 
	let regcounts = argRegs vmc
	    (oldreg,oldcount) = maximumBy argcomp regcounts
	    argcomp ((ArgRegister x),_) ((ArgRegister y),_) = x `compare` y
	    (newreg,newcount) = if regcounts /= []
	                           then (inc oldreg, oldcount+1)
	                           else (ArgRegister 0, 1)
	    inc (ArgRegister x) = (ArgRegister (x + 1))
	    name = if isJust vname
	              then fromJust vname
	              else (anonname newreg)
	in (newreg,insert name (newreg,newcount) vmc)

assignTempRegister ::  VarMapCount -> (Register,VarMapCount)
assignTempRegister  vmc = 
	let regcounts = tempRegsC vmc
	    (oldreg,oldcount) = maximumBy tempcomp regcounts
	    tempcomp ((TempRegister x),_) ((TempRegister y),_) = x `compare` y
	    (newreg,newcount) = if regcounts /= []
	                           then (inc oldreg, oldcount+1)
	                           else (TempRegister 0, 1)
	    inc (TempRegister x) = (TempRegister (x + 1))
	in (newreg,insert (anonname newreg) (newreg,newcount) vmc)

anonname x = "ANONREG" ++ (show x)

argRegs :: VarMapCount -> [(Register,Int)]
argRegs = elems . (Data.Map.filter fil)
           where fil ((ArgRegister x),_) = True
                 fil _ = False
                 
tempRegsC :: VarMapCount -> [(Register,Int)]
tempRegsC = elems . (Data.Map.filter fil)
           where fil ((TempRegister x),_) = True
                 fil _ = False


-- sketch - compileTerm returns a list of instructions, and a function
--          to produce the reference instruction

compileTerm :: Term -> VarMap -> AMap -> [WAMOp]
compileTerm (T _ _ args _) vm0 amap = 
         let (ret,_,_) = foldl (accumulate True amap) ([],[],vm1) args
             vm1 = varMapToVarMapCount vm0
         in ret

compileTermHack :: Term -> VarMap -> AMap -> ([WAMOp],VarMapCount)
compileTermHack (T _ _ args _) vm0 amap = 
         let (ret,_,vmc) = foldl (accumulate True amap) ([],[],vm1) args
             vm1 = varMapToVarMapCount vm0
         in (ret,vmc)

compileTerm1 :: Term -> VarMapCount -> Bool -> AMap -> ([WAMOp],WAMOp,VarMapCount )
compileTerm1 (T at ar args _) vm0 toplevel amap =
    ( [(Get_structure (findWithDefault 0 at amap) ar structreg)] ++
            regassigns 
                 ++ argumentinstrs ,
                 (Unify_variable structreg),
                   vm2)
    where (structreg,vm1) = if toplevel
                               then assignArgRegister Nothing vm0
                               else assignTempRegister  vm0
          (argumentinstrs, regassigns, vm2) =
              foldl (accumulate False amap) ([],[],vm1) args

-- Every variable is guaranteed a global register assignment
compileTerm1 (V v _) vm0 True _ =
	let (norm_reg,_) = (findWithDefault (NullRegister,0) v vm0)
	    (areg,vm1) = assignArgRegister Nothing vm0
	in if not (seenVariable v vm0)
          then ([(Get_variable norm_reg areg)],NullOP,(addSeen v vm1))
          else ([(Get_value norm_reg areg)],NullOP,vm1)

compileTerm1 (V v _) vm False _ =
	let (norm_reg,_) = (findWithDefault (NullRegister,0) v vm)
	in if not (seenVariable v vm)
          then ([],(Unify_variable norm_reg),(addSeen v vm))
          else ([],(Unify_value norm_reg),vm)

accumulate :: Bool -> AMap -> ([WAMOp],[WAMOp],VarMapCount) -> Term -> ([WAMOp],[WAMOp],VarMapCount)
accumulate depth amap (insts,regs,vm0) term =
    let (insts1,reg,vm1) = compileTerm1 term vm0 depth amap
    in ((insts ++ insts1),regs ++ [reg], vm1)


leftmost (CCONJ x y _) = leftmost x
leftmost (CDISJ x y _) = leftmost x
leftmost (CTERM x _) = x

-- a bit hacked, assumes term is nice
deleteLeftmostCut (CCONJ x y v) = (CCONJ (deleteLeftmostCut x) y v)
deleteLeftmostCut (CDISJ x y v) = (CDISJ (deleteLeftmostCut x) y v)
deleteLeftmostCut (CTERM (CutTerm _) _) = CNIL
deleteLeftmostCut x = x
        
containsCut (CCONJ x y _) = (containsCut x) || (containsCut y)
containsCut (CDISJ x y _) = (containsCut x) || (containsCut y)
containsCut (CTERM (CutTerm _) _) = True
containsCut _ = False
        
    
{- TEST FUNCTIONS -}

--
-- Builtins down here for module dependencies ...
--

stopcounts :: WAMData a -> IO ()
stopcounts wd = writeIORef (countsRunning wd) False

startcounts :: WAMData a -> IO ()
startcounts wd = writeIORef (countsRunning wd) True

traceon :: WAMData a -> IO ()
traceon wd = writeIORef (traceall wd) True

traceoff :: WAMData a -> IO ()
traceoff wd = writeIORef (traceall wd) False

builtinVar :: WAMData a -> IO ()
builtinVar wd =
    do ref0 <- getRegister wd (registerToRaw $ ArgRegister 0)
       der <- deref wd (deref1 ref0)
       cont <- getHEntry wd der
       if (isVar der cont)
          then return ()
          else doWAMFail wd

builtinNonVar :: WAMData a -> IO ()
builtinNonVar wd =
    do ref0 <- getRegister wd (registerToRaw $ ArgRegister 0)
       der <- deref wd (deref1 ref0)
       cont <- getHEntry wd der
       if (isVar der cont)
          then doWAMFail wd
          else return ()

builtinTermTest :: (String -> String -> Bool) -> WAMData a -> IO ()
builtinTermTest op wd =
    do ref0 <- getRegister wd (registerToRaw $ ArgRegister 0)
       ref1 <- getRegister wd (registerToRaw $ ArgRegister 1)
       val0 <- deref wd (deref1 $ ref0)
       val1 <- deref wd (deref1 $ ref1)
       cont0 <- getHEntry wd val0
       cont1 <- getHEntry wd val1
       str0 <- Engine.lookupAtom wd (st cont0)
       str1 <- Engine.lookupAtom wd (st cont1)
       dd ("TERM COMPARING " ++ (show val0) ++ " AND " ++ (show val1) ++ "\n")
       if op (fromJust str0) (fromJust str1)
          then return ()
          else doWAMFail wd
    where st hfunc = fst $ getHFunc hfunc

builtinIS :: WAMData a -> IO ()
builtinIS wd =
    do ref0 <- getRegister wd (registerToRaw $ ArgRegister 0)
       der <- deref wd (deref1 ref0)
       cont <- getHEntry wd der
       if not (isVar der cont)
          then fail "First argument of \"is\" must be a variable"
          else return()
       ref1 <- getRegister wd (registerToRaw $ ArgRegister 1)
       val1 <- arithEval wd (deref1 ref1)
       machineval1 <- Engine.newAtom wd (show val1)
       put_structure machineval1 0 itempreg wd
       get_value (registerToRaw $ ArgRegister 0) itempreg wd
       dd ("RESULT OF IS: " ++ (show val1) ++ "\n")
       return ()

builtinArithTest :: (Int64 -> Int64 -> Bool) -> WAMData a -> IO ()
builtinArithTest op wd =
    do ref0 <- getRegister wd (registerToRaw $ ArgRegister 0)
       ref1 <- getRegister wd (registerToRaw $ ArgRegister 1)
       val0 <- arithEval wd (deref1 $ ref0)
       val1 <- arithEval wd (deref1 $ ref1)
       dd ("COMPARING " ++ (show val0) ++ " AND " ++ (show val1) ++ "\n")
       if op val0 val1
          then return ()
          else doWAMFail wd

-- This is a hacked version of IS -- we need unary operators to do this properly
sqrtBI :: WAMData a -> IO ()
sqrtBI wd =
    do ref0 <- getRegister wd (registerToRaw $ ArgRegister 0)
       der <- deref wd (deref1 ref0)
       cont <- getHEntry wd der
       if not (isVar der cont)
          then fail "First argument of \"is\" must be a variable"
          else return()
       ref1 <- getRegister wd (registerToRaw $ ArgRegister 1)
       val1 <- arithEval wd (deref1 ref1)
       machineval1 <- Engine.newAtom wd (show (sqrtfunc val1))
       put_structure machineval1 0 itempreg wd
       get_value (registerToRaw $ ArgRegister 0) itempreg wd
       dd ("RESULT OF IS: " ++ (show val1) ++ "\n")
       return ()
    where sqrtfunc x = fromIntegral $ ceiling $ sqrt $ fromIntegral x


arithop x = x `elem` ["+","-","*","/"]
evbinary "+" a b = a + b
evbinary "*" a b = a * b
evbinary "-" a b = a - b
evbinary "/" a b = a `div` b

unaryop x = x `elem` ["sqrt"]
-- XXX dodgy integer square roots !
evunary "sqrt" x = fromIntegral $ ceiling $ sqrt $ fromIntegral x


{-

-- Version for in-place integers when ready

arithEval :: (WAMData a) -> WAMRef -> IO Int
arithEval wd ref = 
	do entryref <- deref wd ref
	   entry <- getHEntry wd entryref
	   case entry of
		e | (isHInt e) -> return $ derefInt e
		  | (isHFunc e) && (arity == 2) ->
		      do op <- lookupAtom wd rawop
			 if (arithop $ fromJust op)
			    then do val1 <- arithEval wd (entryref+1)
				    val2 <- arithEval wd (entryref+2)
				    return (ev1 (fromJust op) val1 val2)
			    else afail $ " not arithmetic operator " ++ (show op)
		  | otherwise -> afail $ " not arithmetic atom or operator " ++ (pentry entry)
		  where afail x = fail ("Attempting arithmetic with non-arithmetic operator;" ++ x)
			(rawop,arity) = getHFunc entry

-}

arithEval :: (WAMData a) -> WAMRef -> IO Int64
arithEval wd ref = 
	do entryref <- deref wd ref
	   entry <- getHEntry wd entryref
	   if not (isHFunc entry)
	      then afail $ " not arithmetic atom or operator " ++ (pentry entry)
	      else do atstr <- lookupAtom wd (fst $ getHFunc entry)
		      case (fromJust atstr,snd $ getHFunc entry) of
		       (op,2) -> if not (arithop op)
				 then afail $ " not binary arithmetic operator " ++ (show op)
				 else do val1 <- arithEval wd (entryref+1)
					 val2 <- arithEval wd (entryref+2)
					 return $! (evbinary op val1 val2)
		       (op,1) -> if not (unaryop op)
				 then afail $ " not unary arithmetic operator " ++ (show op)
				 else do val1 <- arithEval wd (entryref+1)
					 return $! (evunary op val1)
		       (int,0) -> case  (readSigned readDec int) of
					 [(num,str)] | str == "" -> return num
						     | otherwise -> afail $ " not arithmetic atom " ++ (show int)
		       otherwise -> afail $ " not arithmetic atom " ++ (show atstr)

	where afail x = fail ("Attempting arithmetic with non-arithmetic operator;" ++ x)


getCountsBI :: (WAMData a) -> IO ()
getCountsBI wd =
    do countid <- getRegister wd (registerToRaw $ ArgRegister 0)
       counttype <- getRegister wd (registerToRaw $ ArgRegister 1)
       ref2 <- getRegister wd (registerToRaw $ ArgRegister 2)
       deref2 <- deref wd (deref1 ref2)
       cont2 <- getHEntry wd deref2
       if not (isVar deref2 cont2)
          then fail "Third argument of get_counts must be a variable"
          else return()
       intcid <- arithEval wd (deref1 $ countid)
       intct <- arithEval wd (deref1 $ counttype)
       let lkup = if intct == 0
		     then fst
		     else snd
	   in do ct <- getCounts wd (fromIntegral intcid)
		 machineval1 <- Engine.newAtom wd (show (lkup ct))
		 put_structure machineval1 0 itempreg wd
		 get_value (registerToRaw $ ArgRegister 2) itempreg wd
		 dd ("RESULT OF GET COUNTS: " ++ (show $ lkup ct) ++ "\n")
	         return ()

-- placeholder for above
getCounts :: a -> b -> IO (Int,Int)
getCounts _ _ = return (0,0)

setWeightBI :: (WAMData a) -> IO ()
setWeightBI wd =
    do ref0 <- getRegister wd (registerToRaw $ ArgRegister 0)
       val0 <- arithEval wd (deref1 $ ref0)
       writeIORef (weightMultiple wd) (fromIntegral val0)
       return ()


printTerm :: (WAMData a) -> IO ()
printTerm wd = do ref <- getRegister wd (registerToRaw $ ArgRegister 0)
                  term <- cookedTerm wd (derefdebug 12 ref)
                  putStrLn (show term)
                  
      

printAddr :: (WAMData a) -> WAMRef -> IO ()
printAddr wd ref = do term <- cookedTerm wd ref
                      putStrLn (show term)


termSize :: (WAMData a) -> WAMRef -> IO Int
termSize wd ref = do term <- cookedTerm wd ref
                     return (cts term)

cts (T _ _ a _) = (sum (Data.List.map cts a)) + 1
cts _ = 1

biFail :: (WAMData a) -> IO ()
biFail wd = doWAMFail wd

-- This is a bit hacky
readTermBI :: (WAMData a) -> IO ()
readTermBI wd = do fail "READ TERM IS BROKEN!"
		   ref <- getRegister wd (registerToRaw $ ArgRegister 0)
                   eterm <- readInputTerm
                   either 
                     (\_ -> do putStrLn "Parse Error on Input"
                               doWAMFail wd)
                     (\term  -> do codetop <- getCT wd
                                   oldpc <- getPC wd
                                   -- NOW THORIBROKEN
                                   compileHead wd term (varMap (CLA term CNIL)) Data.Map.empty
                                   pushCode wd HaltOP
                                   cc <- dumpCode wd
                                   putStrLn cc
                                   setPC wd codetop
                                   execCode wd -- XXX this should be sreg
                                   setPC wd oldpc
                                   setCT wd codetop)
                     eterm

type BIEntry a = ((Atom,Arity),Builtin a)
type BuiltinList a = [BIEntry a]

builtinList :: BuiltinList a
builtinList = [(("print",1),(BI printTerm)),
               (("fail",0),(BI biFail)),
               (("read",1),(BI readTermBI)),
               (("var",1),(BI builtinVar)),
               (("nonvar",1),(BI builtinNonVar)),
               ((">",2),(BI (builtinArithTest (>)))),
               (("=<",2),(BI (builtinArithTest (<=)))),
               (("<",2),(BI (builtinArithTest (<)))),
               ((">=",2),(BI (builtinArithTest (>=)))), 
               (("=:=",2),(BI (builtinArithTest (==)))),              
               (("is",2),(BI builtinIS)),
               (("@>",2),(BI (builtinTermTest (>)))),
               (("@<",2),(BI (builtinTermTest (<)))),
               (("@>=",2),(BI (builtinTermTest (>=)))),
               (("@=<",2),(BI (builtinTermTest (<=)))),
	       (("sqrt",2),(BI (sqrtBI ))),
	       (("stopcounts",0),(BI stopcounts)),
	       (("startcounts",0),(BI startcounts)),
	       (("traceon",0),(BI traceon)),
	       (("traceoff",0),(BI traceoff)),
	       (("get_counts",3),(BI getCountsBI)),
	       (("set_weight_multiple",1),(BI setWeightBI))
              ]
                 
builtinListCompile :: (WAMData a) -> BuiltinList a -> IO [((AtomRef,Arity),Builtin a)]
builtinListCompile wd list =
    do lc <- mapM (compileBI wd) list
       return lc

compileBI :: (WAMData a) -> BIEntry a -> IO ((AtomRef,Arity),Builtin a)
compileBI wd ((atstr,ar),bi) =
    do at <- Engine.newAtom wd atstr
       return ((at,ar),bi)

emptyWAM :: UnsignedNum -> UnsignedNum -> UnsignedNum -> IO (WAMData a)
emptyWAM heapsize codesize trailsize = 
    do cwd <- newWAMData
       instc <- newICData
       callc <- newICData
       heap <- (newArray ( 0, (heapsize + argRegSize*2) ) (0))
       code <- (newArray ( 0,  codesize) 0)
       codeaux <- (newArray ( 0,  codesize) 0)
       trail <- (newArray (0,  trailsize) TEmpty )
       callstack <- (newIORef [])
       heaptop <- newIORef ( 0)
       codetop <- newIORef ( 0)
       envtop <- newIORef ( heapsize)
       choice <- newIORef ( heapsize)
       cb <- newIORef ( heapsize)
       cutchoice <- newIORef ( heapsize)
       dedchoice <- newIORef 0
       regc1 <- newIORef ( 0)
       regc2 <- newIORef 0
       regc3 <- newIORef 0
       sreg <- newIORef 0
       tt <- newIORef 0
       pc <- newIORef ( 0)
       cp <- newIORef (0)
       point <- newIORef 0
       ci <- newIORef Data.Map.empty
       fm <- newIORef Data.Map.empty
       bi <- newIORef Data.Map.empty
       m <- newIORef Read
       q <- newIORef True
       oct <- newIORef 0
       rt <- newIORef (emptyTable 0 lambdaDefault weightScalePowerDefault)
       rng <- newIORef (mkStdGen randSeedDefault)
       rp <- newIORef (Set.empty)
       crun <- newIORef True
       traceall <- newIORef False
       och <- newIORef 0
       -- lam <- newIORef lambdaDefault
       wm <- newIORef 1
       ai <- newIORef (AtomMap.empty)
       cix <- newIORef (AtomMap.empty)
       ctc <- newIORef trialCount
       itc <- newIORef indTrialCountDefault
       rth <- newIORef restartThresholdDefault
       sth <- newIORef shortThresholdDefault
       seed <- newIORef randSeedDefault
       iq <- newIORef CNIL
       dq <- newIORef CNIL
       let (nullCIDRef,ci0) = AtomMap.newAtom NullCID (AtomMap.empty)
	   (exceptCIDRef,ci1) = AtomMap.newAtom ExceptCID ci0
           iwd = (WT {cWAMData = cwd,
                      cInstCount = instc,
                      cCallCount = callc,
                      heap = heap,
                      code = code,
		      codeaux = codeaux,
                      trail = trail,
                      callstack = callstack,
                      heapTop = heaptop,
                      codeTop = codetop,
                      envTop = envtop,
                      choice = choice,
                      choiceBottom = cb,
                      dedchoice = dedchoice,
                      cutchoice = cutchoice,
                      variableMap = fm,
                      heapSize =   heapsize,
                      trailSize =  trailsize,
                      codeSize =  codesize,
                      trailTop = tt,
                      argreg = regc1,
                      tempreg = regc2,
                      permreg = regc3,
                      callIndex = ci,
                      builtinIndex = bi,
                      queryResult = q,
                      pc = pc,
                      oldCodeTop = oct,
                      sreg = sreg,
                      contp = cp,
                      randTable = rt,
		      -- lambdaVal = lam,
                      mode = m,
                      sgen = rng,
                      randpreds = rp,
		      countsRunning = crun,
		      traceall = traceall,
		      atomIndex = ai,
		      clauseIDIndex = cix,
		      weightMultiple = wm,
		      nullCID = nullCIDRef,
		      exceptCID = exceptCIDRef,
                      indTrialCount = itc,
		      EngineTypes.restartThreshold = rth,
		      EngineTypes.shortThreshold = sth,
		      pointChange = point,
		      origChange = och,
		      seed = seed,
		      currTrialCount = ctc,
                      dedQuery = dq,
                      indQuery = iq})
	   in do bil <- builtinListCompile iwd builtinList
		 bie <- newIORef (fromList bil)
		 cix <- newIORef (ci1)
		 return (iwd { builtinIndex = bie,clauseIDIndex = cix })

